import io
import os
import sys

from flask import Flask, request, send_file, jsonify
from flask_cors import CORS
from PIL import Image
import numpy as np
import time
import logging
import u2net
import uuid
import base64
import re

logging.basicConfig(level=logging.INFO)

# Initialize the Flask application
app = Flask(__name__)
CORS(app)

@app.route('/api/hello', methods=['GET'])
def hello():
    return '<h1>Background Removal API</h1>'


# Route http posts to this method
@app.route('/api/removebg', methods=['POST'])
def removebg():
    start = time.time()
    img = base64_to_image(request.form.get("base64Data"));
    # Convert string data to PIL Image
    # img = Image.open(io.BytesIO(data))
    # Ensure i,qge size is under 1024
    if img.size[0] > 1024 or img.size[1] > 1024:
        img.thumbnail((1024, 1024))

    # Process Image
    mask = u2net.run(np.array(img)).convert("L")
    bgColor = (255,255,255)
    if (request.form.get("backgroundColor") == '蓝色'):
        bgColor = (71, 151, 221)
    elif (request.form.get("backgroundColor") == '红色'):
        bgColor = (200, 44, 59)

    img_bg_removed = naive_cutout(img, mask, bgColor)
    img_path = save_img(img_bg_removed)
    base64Data = image_to_base64(img_bg_removed)
    logging.info(f'Completed in {time.time() - start:.2f}s')
    return jsonify({'code': 200, 'base64': base64Data, 'imgPath': "http://localhost:8081/api/downloadImg?imgName=" + img_path})

@app.route('/api/downloadImg', methods=['GET'])
def download_img():
    imgName = request.args.get('imgName')
    full_path = os.path.join(os.getcwd(), 'output_img', imgName)
    img = Image.open(full_path)
    buff = io.BytesIO()
    img.save(buff, 'JPEG')
    buff.seek(0)
    return send_file(buff, mimetype='image/jpeg');

def base64_to_image(base64_str):
    base64_data = re.sub('^data:image/.+;base64,', '', base64_str)
    byte_data = base64.b64decode(base64_data)
    image_data = io.BytesIO(byte_data)
    img = Image.open(image_data)
    return img

def image_to_base64(img):
    # Save to buffer
    buff = io.BytesIO()
    img.save(buff, 'JPEG')
    buff.seek(0)
    imgData = "data:image/jpeg;base64," + str(base64.b64encode(buff.getvalue())).replace("b'", "").replace("'", "");
    return imgData;

def save_img(img):
    output_img = os.path.join(os.getcwd(), 'output_img')
    if not os.path.exists(output_img):
        os.makedirs(output_img, exist_ok=True)
    ima_name = str(uuid.uuid4()) + ".png"
    img_path = output_img + os.sep + ima_name
    img.save(img_path)
    return ima_name

def naive_cutout(img, mask, backgroundColor=(255, 255, 255)):
    empty = Image.new("RGB", (img.size), backgroundColor)
    cutout = Image.composite(img, empty, mask.resize(img.size, Image.LANCZOS))
    return cutout

if __name__ == '__main__':
    os.environ['FLASK_ENV'] = 'development'
    port = int(os.environ.get('PORT', 8081))
    app.run(debug=True, host='localhost', port=port)